import logging
from sqlalchemy.orm import Session

from app.db.session import Base, engine, SessionLocal
from app.models import MCPCategory, MCPTool

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def init_db() -> None:
    """
    初始化数据库表和默认数据
    """
    # 创建所有表
    Base.metadata.create_all(bind=engine)
    logger.info("Created database tables")


def create_initial_data(db: Session) -> None:
    """
    创建初始测试数据
    """
    # 检查分类是否已存在
    if db.query(MCPCategory).count() > 0:
        logger.info("Initial data already exists")
        return

    # 创建默认分类
    image_generation = MCPCategory(
        name="图像生成",
        description="AI图像生成工具",
        icon="image"
    )
    
    search = MCPCategory(
        name="搜索工具",
        description="各类搜索工具",
        icon="search"
    )
    
    nlp = MCPCategory(
        name="自然语言处理",
        description="文本分析和处理工具",
        icon="text"
    )
    
    db.add_all([image_generation, search, nlp])
    db.commit()
    
    # 创建默认工具
    tools = [
        MCPTool(
            name="文生图",
            english_name="text2image",
            description="根据文本描述生成图像",
            icon="/icons/text2image.png",
            category_id=image_generation.id,
            api_endpoint="/api/v1/tools/text2image",
            parameters_schema={
                "properties": {
                    "prompt": {
                        "type": "string",
                        "description": "图像描述"
                    },
                    "size": {
                        "type": "string",
                        "enum": ["256x256", "512x512", "1024x1024"],
                        "default": "512x512"
                    },
                    "n": {
                        "type": "integer",
                        "minimum": 1,
                        "maximum": 4,
                        "default": 1
                    }
                },
                "required": ["prompt"]
            }
        ),
        MCPTool(
            name="联网搜索",
            english_name="web_search",
            description="搜索互联网获取信息",
            icon="/icons/search.png",
            category_id=search.id,
            api_endpoint="/api/v1/tools/web_search",
            parameters_schema={
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "搜索查询"
                    },
                    "limit": {
                        "type": "integer",
                        "default": 5,
                        "minimum": 1,
                        "maximum": 10
                    }
                },
                "required": ["query"]
            }
        ),
        MCPTool(
            name="文本分类",
            english_name="text_classification",
            description="对文本进行分类",
            icon="/icons/classify.png",
            category_id=nlp.id,
            api_endpoint="/api/v1/tools/classify",
            parameters_schema={
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "要分类的文本"
                    },
                    "categories": {
                        "type": "array",
                        "items": {
                            "type": "string"
                        },
                        "description": "可选的分类类别"
                    }
                },
                "required": ["text"]
            }
        )
    ]
    
    db.add_all(tools)
    db.commit()
    
    logger.info("Created initial data")


if __name__ == "__main__":
    logger.info("Initializing database...")
    init_db()
    
    # 创建会话并添加初始数据
    db = SessionLocal()
    try:
        logger.info("Creating initial data...")
        create_initial_data(db)
        logger.info("Database initialization completed successfully")
    finally:
        db.close() 